"""
@Filename       : dnrl.py
@Create Time    : 2021/1/12 11:08
@Author         : Rylynn
@Description    : Joint Learning of User Representation with Diffusion Sequence and Network Structure
                TKDE 2020
"""

import torch as th
import torch.nn as nn


class NDRL(nn.Module):
    def __init__(self, config):
        super(NDRL, self).__init__()
        self.sender_embed = nn.Embedding(config['user_num'], config['embed_size'])
        self.receiver_embed = nn.Embedding(config['user_num'], config['embed_size'])
        self.scale = config['scale']
        self.beta = config['beta']

    def influence_func(self, senders, receivers):
        return self.scale * th.sigmoid(self.beta * th.norm(senders - receivers, 2) / 2)

    def infection_func(self, senders, receivers, time_from, time_end):
        influence = self.influence_func(senders, receivers)
        return influence * th.exp(- influence * (time_end - time_from))

    def survival_func(self, senders, receivers, time_from, time_end):
        influence = self.influence_func(senders, receivers)
        return th.exp(- influence * (time_end - time_from))

    def forward(self,):
        pass

    def loss(self):
        pass

